import torch

probs = torch.tensor([[0.1, 0.2, 0.7], [0.5, 0.3, 0.2]])  # 简化的概率分布，形状(2,3)
token_ids = torch.multinomial(probs, num_samples=1).squeeze(-1)
print(token_ids)


import torch

# 定义参数（与你的场景一致）
block_num = 2341  # 块数量
qk_rope_head_dim = 128  # RoPE头维度
KVCACHE_NZ_DIM = 16  # 缓存分块维度
block_size = 128  # 每个块的token数量

# 1. 模拟原始kv_cache[1]的张量（形状：[2341, 128, 1, 64]）
# 注意：原形状总元素数需与目标形状匹配，这里修正原形状最后一维为128以满足匹配
# （原问题中64可能为笔误，否则总元素数不匹配，无法view）
kv_cache_1 = torch.randn(2341, 128, 1, 128)  # 随机生成模拟数据
print(f"原始kv_cache[1]形状: {kv_cache_1.shape}")
print(f"原始总元素数: {kv_cache_1.numel()}")

# 2. 计算中间维度（qk_rope_head_dim // KVCACHE_NZ_DIM）
split_dim = qk_rope_head_dim // KVCACHE_NZ_DIM
print(f"\n拆分维度 (qk_rope_head_dim // KVCACHE_NZ_DIM): {split_dim}")  # 输出：8

# 3. 执行view重塑操作
kv_cache_pe = kv_cache_1.view(
    block_num,
    1,
    split_dim,
    block_size,
    KVCACHE_NZ_DIM
)

# 4. 验证结果
print(f"\n重塑后kv_cache_pe形状: {kv_cache_pe.shape}")
print(f"重塑后总元素数: {kv_cache_pe.numel()}")

# 5. 验证总元素数是否匹配（view操作要求前后元素数必须相等）
assert kv_cache_1.numel() == kv_cache_pe.numel(), "元素数不匹配，view操作失败！"
print("\n验证成功：重塑前后总元素数一致")
